Skip to content

feat(kernel): thread Azure Entra OAuth (U2M + SP M2M) through the auth bridge - #919

Open
eric-wang-1990 wants to merge 4 commits into
mainfrom
eric-wang/kernel-azure-oauth
Open

feat(kernel): thread Azure Entra OAuth (U2M + SP M2M) through the auth bridge#919
eric-wang-1990 wants to merge 4 commits into
mainfrom
eric-wang/kernel-azure-oauth

Conversation

@eric-wang-1990

Copy link
Copy Markdown
Contributor

What

Teaches the use_kernel=True auth bridge to support Azure Entra (Azure AD) OAuth — both the U2M browser flow (azure-oauth) and the service-principal M2M flow (azure-sp-m2m). Previously azure-oauth was rejected with NotSupportedError and azure-sp-m2m fell through to the generic "unsupported" error.

The kernel needs no Azure-specific code: both types route onto the kernel's generic OAuth flows with Azure values supplied as overrides. This depends on the kernel-side override plumbing added in databricks/databricks-sql-kernel#263 (token_url / scope overrides on M2M + U2M; M2M skips OIDC discovery when a token endpoint is pinned).

How

src/databricks/sql/backend/kernel/auth_bridge.py:

  • azure-oauth (U2M) → kernel oauth-u2m with the Azure app client id (96eecda7-…), redirect port 8030, and the AAD delegated scope {app_id}/user_impersonation offline_access (synthesised via AzureOAuthEndpointCollection.get_scopes_mapping, honoring DATABRICKS_AZURE_TENANT_ID). The kernel runs the browser flow and discovers endpoints via the workspace /oidc redirector (which an Azure workspace redirects to Entra). PECOBLR-4120.
  • azure-sp-m2m (M2M) → kernel oauth-m2m with the Azure service-principal credentials, an Entra v2.0 token_url (login.microsoftonline.com/{tenant}/oauth2/v2.0/token), and the {effective_app_id}/.default scope. PECOBLR-4141.
    • Requires an explicit azure_tenant_id (the kernel path does not auto-discover the tenant from the workspace's /aad/auth redirect as the Thrift path does) — raises a clear error otherwise.
    • The Azure management-token header (X-Databricks-Azure-SP-Management-Token) and azure_workspace_resource_id are not applied on the kernel path. No SQL connector (Go, Node) uses them; the Databricks-audience token authenticates SPs that are workspace principals (the SQL norm). An RBAC-only-not-in-workspace SP is unsupported here (documented; warns if azure_workspace_resource_id is set).

kernel_auth_kwargs now takes hostname (for the effective Azure app id per cloud); client.py passes self._server_hostname.

Tests

TDD — replaced the azure-oauth NotSupportedError test with routing tests, and added a TestKernelAzureSpM2M suite (routing, required tenant, required creds, federation client id). 50 bridge tests pass; black clean. (test_kernel_client.py skips without the kernel wheel, as before.)

Related

  • Kernel side: databricks/databricks-sql-kernel#263 (design + implementation).
  • JIRA: PECOBLR-4141 (azure-sp-m2m), PECOBLR-4120 (azure-oauth U2M).

This pull request and its description were written by Isaac.


This PR was created with GitHub MCP.

…h bridge

The kernel auth bridge rejected azure-oauth and had no azure-sp-m2m path.
Route both Azure auth types onto the kernel's generic OAuth flows (the
kernel needs no Azure-specific code; PR databricks/databricks-sql-kernel#263
added the token_url/scope override plumbing this relies on):

- azure-oauth (Azure AD U2M) -> oauth-u2m with the Azure app client id
  (96eecda7-...), redirect port 8030, and the {app_id}/user_impersonation
  offline_access delegated scope (via AzureOAuthEndpointCollection, honoring
  DATABRICKS_AZURE_TENANT_ID). The kernel discovers endpoints via the
  workspace /oidc redirector. (PECOBLR-4120)
- azure-sp-m2m (Azure service principal) -> oauth-m2m with the Azure creds,
  an Entra v2.0 token_url, and the {effective_app_id}/.default scope.
  Requires an explicit azure_tenant_id (the kernel path does not
  auto-discover it). The management-token header / azure_workspace_resource_id
  are not applied on the kernel path -- no SQL connector uses them, matching
  Go and Node. (PECOBLR-4141)

kernel_auth_kwargs now takes hostname (for the effective Azure app id);
the client passes self._server_hostname.

TDD: replaced the azure-oauth NotSupportedError test with routing tests and
added a TestKernelAzureSpM2M suite (routing, required tenant/creds,
federation client id). 50 bridge tests pass; black clean.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR extends the use_kernel=True auth bridge to support Azure Entra (Azure AD) OAuth by routing the connector’s Azure auth types (azure-oauth U2M and azure-sp-m2m M2M) onto the kernel’s generic OAuth flows with Azure-specific overrides.

Changes:

  • Route auth_type="azure-oauth" to kernel oauth-u2m using the Azure OAuth app bundle and Azure scope mapping.
  • Route auth_type="azure-sp-m2m" to kernel oauth-m2m using Entra v2.0 token URL and {effective_app_id}/.default scope; require azure_tenant_id.
  • Thread hostname into kernel_auth_kwargs (via KernelDatabricksClient.open_session) and add/adjust unit tests + changelog entry.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File Description
src/databricks/sql/backend/kernel/auth_bridge.py Adds Azure auth-type routing and related override plumbing for kernel session kwargs.
src/databricks/sql/backend/kernel/client.py Passes the server hostname into kernel_auth_kwargs during session open.
tests/unit/test_kernel_auth_bridge.py Updates U2M tests for azure-oauth routing and adds azure-sp-m2m routing/validation tests.
CHANGELOG.md Documents Azure Entra OAuth support on the kernel backend.
Suppressed comments (1)

src/databricks/sql/backend/kernel/auth_bridge.py:262

  • auth_type='azure-sp-m2m' uses azure_client_id/azure_client_secret, but because this branch returns before the generic ambiguity guards, a caller can also supply oauth_client_id/oauth_client_secret and have them silently ignored. This makes misconfigurations hard to diagnose.
    if auth_type == "azure-sp-m2m":
        azure_client_id = opts.get("azure_client_id")
        azure_client_secret = opts.get("azure_client_secret")
        azure_tenant_id = opts.get("azure_tenant_id")
        if not (azure_client_id and azure_client_secret):

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +227 to +229
if auth_type == "azure-oauth":
raise NotSupportedError(
"use_kernel=True does not support auth_type='azure-oauth' (Azure "
"AD U2M) yet: the kernel resolves OAuth endpoints only from the "
"workspace-native OIDC configuration and cannot drive the Azure AD "
"authorization/token flow. Use the Thrift backend (default) for "
"azure-oauth. Tracked by PECOBLR-4120."
redirect_port = opts.get("oauth_redirect_port")
caller_scopes = _normalize_scopes(opts.get("oauth_scopes"))
Comment on lines +284 to +285
app_id = get_effective_azure_login_app_id(hostname or "")
kwargs = {

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 Medium

Routing logic is clean and well-tested overall. One medium concern: the new Azure branches return before the ambiguity guards, so azure-oauth + oauth_client_secret (or + credentials_provider) is silently accepted as U2M whereas the parallel databricks-oauth case raises NotSupportedError. Nit (not filed inline): the step-3 comment # Only databricks-oauth reaches here (azure-oauth rejected up front) in the U2M block is now stale — azure-oauth is routed up front, not rejected.

# before the generic M2M/PAT/U2M routing below (azure-sp-m2m carries its
# creds in azure_* kwargs, not oauth_client_id/secret, so it would
# otherwise fall through to the final "unsupported" error).

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium — The azure-oauth (and azure-sp-m2m) branches return before the ambiguity guards run, so conflicting signals that are rejected for the non-Azure U2M path are silently accepted for Azure.

Concretely, auth_type="azure-oauth" + oauth_client_secret reaches this branch, builds oauth-u2m kwargs, and returns — the secret is silently discarded and the browser flow runs. The exactly-parallel non-Azure case (auth_type="databricks-oauth" + oauth_client_secret) instead raises NotSupportedError("Ambiguous auth ...") a few lines below. The same gap lets azure-oauth + a custom credentials_provider bypass the credentials_provider-ambiguity guard. This is also a behavior regression: the prior code rejected any azure-oauth request up front, so a mis-specified azure-oauth + secret used to fail loudly and now succeeds with a different principal/flow than the caller likely intended.

Consider applying the same ambiguity checks (secret-with-U2M, credentials_provider-with-secret) to the Azure U2M branch before returning, so an ambiguous Azure request fails at session-open rather than silently choosing the browser flow.

The auth table marked the azure_* fields as Kernel-unsupported and
claimed azure-oauth 'still works on the kernel' (it was actually
rejected). Reflect the new routing: azure-sp-m2m + azure-oauth now work
on the kernel path; azure_tenant_id is required there; the management
token / azure_workspace_resource_id are not applied (matching Go/Node).

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 Low

Looks good — clean, well-documented routing with strong test coverage for both new Azure flows (routing, required tenant/creds, federation client id). One low-severity consistency note: the azure-oauth branch returns before the ambiguity guards, so oauth_client_secret / credentials_provider are silently ignored there, unlike the databricks-oauth U2M path which rejects them.

# workspace /oidc redirector (which an Azure workspace redirects to Entra).
# The AAD delegated scope ({app_id}/user_impersonation [+ offline_access])
# is synthesised via AzureOAuthEndpointCollection, which also honors the
# DATABRICKS_AZURE_TENANT_ID app-id override. PECOBLR-4120.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — The azure-oauth branch returns before the section-0 ambiguity guards, so it silently ignores auth signals that the databricks-oauth U2M path rejects loudly. Specifically:

  • auth_type="azure-oauth" + oauth_client_secret → routes to U2M and drops the secret, whereas auth_type="databricks-oauth" + oauth_client_secret raises NotSupportedError("Ambiguous auth …").
  • auth_type="azure-oauth" + credentials_provider → also silently ignored.

These combinations are nonsensical (Azure M2M is expressed via azure-sp-m2m + azure_client_secret, not oauth_client_secret), so the practical blast radius is small. But for consistency with the U2M guard — and to avoid a confusing silent browser-flow launch when a caller clearly intended a non-interactive credential — consider applying the same ambiguity check to azure-oauth, or documenting the intentional divergence. Not a correctness bug, just an inconsistency in the fail-loud contract.

The kernel is the auth core now: for azure-oauth the bridge forwards only
auth_type='azure-oauth' (+ optional client_id/redirect_port passthrough),
and the kernel pins the workspace v2.0 authorize/token endpoints, the
Azure app client id, port 8030, and the user_impersonation scope. Drops
the connector-side endpoint/scope construction (and the
AzureOAuthEndpointCollection / PYSQL_OAUTH_AZURE_* imports) from the
kernel path. Live-verified end-to-end against an Azure workspace.

azure-sp-m2m still routes to oauth-m2m here pending the kernel's
dedicated azure-sp-m2m variant.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 Medium · 2 Low

Overall solid — the azure-sp-m2m v2.0 .default scope + v2.0 token endpoint correctly mirror the Thrift path's v1.0 resource= form, and the empty-hostname fallback is safe. Main issue is a stale module docstring (F1, medium) that claims the azure-oauth bridge synthesises the client id/port/scope via AzureOAuthEndpointCollection when the code actually just forwards the azure-oauth selector and lets the kernel own resolution — the inline comment already contradicts it. Plus two low doc/consistency nits.

kernel's *generic* OAuth flows with Azure values as overrides (the
kernel needs no Azure-specific code):

- ``azure-oauth`` (U2M) → ``oauth-u2m`` with the Azure app client id

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium — The module docstring here describes azure-oauth as routing → ``oauth-u2m`` with the Azure app client id (``96eecda7-…``), redirect port ``8030``, and the AAD delegated scope ``{app_id}/user_impersonation offline_access`` (via ``AzureOAuthEndpointCollection``…). That does not match what the code actually does. The azure-oauth branch (line 226) forwards auth_type="azure-oauth" (not oauth-u2m) and constructs none of those values — the inline comment at lines 218–225 explicitly states "this binding does NOT construct endpoints/scopes — it just passes auth_type='azure-oauth' plus any optional client_id / redirect_port passthrough" and that the kernel pins the client id / port 8030 / scope.

So the top-of-file docstring directly contradicts both the executable code and the inline comment beside it: it claims the bridge maps to oauth-u2m and synthesises the Azure bundle, when in fact it delegates all of that to the kernel via the azure-oauth selector. A maintainer reading the module docstring will expect AzureOAuthEndpointCollection.get_scopes_mapping to be called here (it never is on this path). Please align the docstring with the actual selector-forwarding behavior.

# a malformed Entra URL), the Azure app client id, port 8030, and the
# `{app_id}/user_impersonation offline_access` scope. So this binding does
# NOT construct endpoints/scopes — it just passes `auth_type='azure-oauth'`
# plus any optional client_id / redirect_port passthrough. PECOBLR-4120.

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — Asymmetric handling of auth_type + oauth_client_secret. For databricks-oauth, supplying oauth_client_secret (a U2M selector plus an M2M secret) is deliberately rejected as "Ambiguous auth" so it doesn't silently authenticate against the wrong principal. But the azure-oauth branch runs before those guards and simply ignores any oauth_client_secret present, silently proceeding with the U2M browser flow. Given azure-oauth is newly supported here, the same class of conflicting-signal request should arguably fail loudly for consistency rather than be silently dropped.

(Anchored to the nearest changed line — see the description for the exact location.)

kernel can't drive the Azure AD authorization/token flow — so it is
rejected up front (PECOBLR-4120)."""
``azure-oauth`` (Azure AD U2M) routes here too — see
``test_azure_oauth_routes_to_kernel_u2m`` (PECOBLR-4120)."""

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — The TestKernelOAuthU2M class docstring points readers to test_azure_oauth_routes_to_kernel_u2m, but no such test exists. The azure-oauth tests are actually named test_azure_oauth_forwards_selector_kernel_owns_resolution and test_azure_oauth_honors_custom_client_id_and_port. Update the cross-reference so it names a real test.

…optional mgmt token

Make the `azure-sp-m2m` bridge thin, matching the kernel becoming the
Azure-aware auth core. The connector now forwards
`auth_type='azure-sp-m2m'` + `azure_client_id` / `azure_client_secret`
(and optional `azure_tenant_id` / `azure_workspace_resource_id`) straight
to the kernel Session, instead of constructing the Entra token endpoint
and `{app_id}/.default` scope itself.

Behavior changes on the kernel path (Thrift parity):
- `azure_tenant_id` is now OPTIONAL — the kernel auto-discovers the
  tenant from the workspace's `/aad/auth` redirect when omitted, exactly
  as the Thrift backend does. (Previously the kernel path required it.)
- `azure_workspace_resource_id` is now honored as an optional add-on:
  forward it and the kernel fetches an Azure-management token and sends
  the `X-Databricks-Azure-SP-Management-Token` +
  `X-Databricks-Azure-Workspace-Resource-Id` pair, so an SP with only an
  Azure RBAC role (not a workspace member) can authenticate. (Previously
  it was dropped with a warning.)

Also thread the `azure_*` connection kwargs into `kernel_auth_options`
in session.py — without this the bridge never received them and
`azure-sp-m2m` failed at session-open with "requires azure_client_id".
Adds a regression test for that threading, and rewrites the bridge tests
for thin forwarding (tenant optional, resource id forwarded). Drops the
now-unused `get_effective_azure_login_app_id` import and
`_AZURE_AAD_LOGIN_HOST` constant.

Co-authored-by: Isaac
Signed-off-by: eric-wang-1990 <e.wang@databricks.com>

@peco-review-bot peco-review-bot Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Verdict: 1 Medium · 1 Low

Looks good overall — the Azure routing logic is thin, well-tested, and correct. Two cleanup items: an unused hostname parameter now threaded through client.py for no effect (medium), and a dead test import (low). Nit: the comment at auth_bridge.py:318 ("azure-oauth rejected up front") is now stale — azure-oauth is routed up front, not rejected.

def kernel_auth_kwargs(
auth_provider: Optional[AuthProvider],
auth_options: Optional[Dict[str, Any]] = None,
hostname: Optional[str] = None,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Medium — The new hostname parameter is never used. kernel_auth_kwargs now accepts hostname (line 152) and client.py was changed to pass self._server_hostname, but the function body never references hostname on any code path — the azure-oauth and azure-sp-m2m branches forward the selector/creds verbatim and let the kernel own resolution.

The PR description says the param exists "for the effective Azure app id per cloud," which matches the intent behind the test-file import of get_effective_azure_login_app_id — but that computation was ultimately delegated to the kernel, leaving this parameter (and the client.py threading) as dead plumbing. Either drop the parameter and the client.py call-site change, or use it. Threading a value in that the function ignores invites future readers to assume it affects resolution when it does not.

PYSQL_OAUTH_SCOPES,
PYSQL_OAUTH_REDIRECT_PORT_RANGE,
)
from databricks.sql.auth.common import get_effective_azure_login_app_id

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🔵 Low — get_effective_azure_login_app_id is imported here but never referenced anywhere in the test module (the TestKernelAzureSpM2M suite asserts only on the thin-forwarding kwargs, which don't involve the effective app id). This is an unused import — flake8/linters will flag it, and it's a leftover from an earlier design where the bridge computed the app id itself. Remove the import.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants